/* 
 * author: fotg
 * @github: https://github.com/ForeGotto
 * @email: fotg@foxmail.com
 */

#pragma region headers
#include <tomcrypt.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <time.h>
#include <stdio.h>
#include <stdlib.h>
#pragma endregion

#pragma region variable declaration
unsigned char plaintext[512], ciphertext[512];
unsigned char tmpkey[512], key[MAXBLOCKSIZE], IV[MAXBLOCKSIZE];
unsigned char inbuf[512], outbuf[512]; /* i/o block size */
unsigned long outlen, y, ivsize, x, decrypt;
int cipher_idx, hash_idx, ks;
char *infile, *outfile, *cipher, *mode;
prng_state prng;
FILE *fdin, *fdout;
int err;
struct stat file_info;
long long octets;
symmetric_ECB ecb;
symmetric_CBC cbc;
symmetric_CFB cfb;
symmetric_OFB ofb;
#pragma endregion

#pragma region tool function
int usage(char *name)
{
#pragma region usage intellisense
  int arr[3] = {12, 0, 4};
  printf("Usage encrypt: %s cipher mode infile outfile\n", name);
  printf("Usage decrypt: %s -d cipher mode infile outfile\n", name);
  printf("Avaliable symmetric cipher: des, aes, rc4, rc6.\n");
  printf("Avaliable symmetric cipher mode: ecb, cbc, cfb, ofb.\n");
  printf("As rc4 works without a mode, just put any word you like as a mode");
  exit(1);
#pragma endregion
}

void ecbmode()
{
#pragma region decryption
  if (decrypt)
  {

#pragma region initialization of algorithm

#pragma region read iv
    /* we need to read iv */
    if (fread(IV, 1, ivsize, fdin) != ivsize)
    {
      printf("Error reading IV from input.\n");
      exit(-1);
    }
#pragma endregion

#pragma region start cipher
    if ((err = ecb_start(cipher_idx, key, ks, 0, &ecb)) != CRYPT_OK)
    {
      printf("ecb_start error: %s\n", error_to_string(err));
      exit(-1);
    }
#pragma endregion

#pragma region figuer size of plain text in octets
    /* we need to remove padding so that the plaintext is the same as origin */
    octets /= ivsize;
#pragma endregion

#pragma endregion

#pragma region decryption loop
    for (long long i = 1; i < octets; i++)
    {
      y = fread(inbuf, 1, ivsize, fdin);

      if ((err = ecb_decrypt(inbuf, plaintext, y, &ecb)) != CRYPT_OK)
      {
        printf("ecb_decrypt error: %s\n", error_to_string(err));
        exit(-1);
      }

#pragma region remove padding
      if (i == octets - 1)
      {
        y = plaintext[ivsize - 1]; //此处ivsize -1需改为密码块长-1
      }
#pragma endregion

      if (fwrite(plaintext, 1, y, fdout) != y)
      {
        printf("Error writing to file.\n");
        exit(-1);
      }
    }
#pragma endregion

#pragma region release file handles
    fclose(fdin);
    fclose(fdout);
#pragma endregion
  }
#pragma endregion

  else

#pragma region encryption
  { /* encrypt */
  /* Setup yarrow for random bytes for IV */

#pragma region initialization of algorithm

#pragma region write iv
    if ((err = rng_make_prng(128, find_prng("yarrow"), &prng, NULL)) != CRYPT_OK)
    {
      printf("Error setting up PRNG, %s\n", error_to_string(err));
    }

    /* You can use rng_get_bytes on platforms that support it */
    /* x = rng_get_bytes(IV,ivsize,NULL);*/
    x = yarrow_read(IV, ivsize, &prng);
    if (x != ivsize)
    {
      printf("Error reading PRNG for IV required.\n");
      exit(-1);
    }

    if (fwrite(IV, 1, ivsize, fdout) != ivsize)
    {
      printf("Error writing IV to output.\n");
      exit(-1);
    }
#pragma endregion

#pragma region start cipher
    if ((err = ecb_start(cipher_idx, key, ks, 0, &ecb)) != CRYPT_OK)
    {
      printf("ecb_start error: %s\n", error_to_string(err));
      exit(-1);
    }
#pragma endregion

#pragma region figuer size of cipher text in octets
    /* we need to pad white space and amount of bytes we pad
        so that the plaintext is multiple of block length
      */
    octets /= ivsize;
    octets += 1;
#pragma endregion

#pragma endregion

#pragma region encryption loop
    for (long long i = 0; i < octets; i++)
    {
      y = fread(inbuf, 1, ivsize, fdin);

#pragma region padding
      if (i == octets - 1)
      {
        for (int cursor = y; cursor < sizeof(inbuf); cursor++)
        {
          inbuf[cursor] = ' ';
        }
        inbuf[ivsize - 1] = (unsigned char)y;
        y = ivsize;
        // printf("encrypt%s\n", inbuf);
      }
#pragma endregion

      if ((err = ecb_encrypt(inbuf, ciphertext, y, &ecb)) != CRYPT_OK)
      {
        printf("ecb_encrypt error: %s\n", error_to_string(err));
        exit(-1);
      }

      if (fwrite(ciphertext, 1, y, fdout) != y)
      {
        printf("Error writing to output.\n");
        exit(-1);
      }
    }
#pragma endregion

#pragma region release file handles
    fclose(fdout);
    fclose(fdin);
#pragma endregion
  }
#pragma endregion
}

void cbcmode()
{
#pragma region decryption
  if (decrypt)
  {

#pragma region initialization of algorithm

#pragma region read iv
    /* we need to read iv */
    if (fread(IV, 1, ivsize, fdin) != ivsize)
    {
      printf("Error reading IV from input.\n");
      exit(-1);
    }
#pragma endregion

#pragma region figuer size of plain text in octets
    /* we need to remove padding so that the plaintext is the same as origin */
    octets /= ivsize;
#pragma endregion

#pragma region start cipher
    if ((err = cbc_start(cipher_idx, IV, key, ks, 0, &cbc)) != CRYPT_OK)
    {
      printf("cbc_start error: %s\n", error_to_string(err));
      exit(-1);
    }
#pragma endregion

#pragma endregion

#pragma region decryption loop
    for (long long i = 1; i < octets; i++)
    {
      y = fread(inbuf, 1, ivsize, fdin);

      if ((err = cbc_decrypt(inbuf, plaintext, y, &cbc)) != CRYPT_OK)
      {
        printf("cbc_decrypt error: %s\n", error_to_string(err));
        exit(-1);
      }

#pragma region remove padding
      if (i == octets - 1)
      {
        y = plaintext[ivsize - 1]; //此处ivsize -1需改为密码块长-1
      }
#pragma endregion

      if (fwrite(plaintext, 1, y, fdout) != y)
      {
        printf("Error writing to file.\n");
        exit(-1);
      }
    }
#pragma endregion

#pragma region release file handles
    fclose(fdin);
    fclose(fdout);
#pragma endregion
  }
#pragma endregion

  else

#pragma region encryption
  { /* encrypt */
  /* Setup yarrow for random bytes for IV */

#pragma region initialization of algorithm

#pragma region write iv
    if ((err = rng_make_prng(128, find_prng("yarrow"), &prng, NULL)) != CRYPT_OK)
    {
      printf("Error setting up PRNG, %s\n", error_to_string(err));
    }

    /* You can use rng_get_bytes on platforms that support it */
    /* x = rng_get_bytes(IV,ivsize,NULL);*/
    x = yarrow_read(IV, ivsize, &prng);
    if (x != ivsize)
    {
      printf("Error reading PRNG for IV required.\n");
      exit(-1);
    }

    if (fwrite(IV, 1, ivsize, fdout) != ivsize)
    {
      printf("Error writing IV to output.\n");
      exit(-1);
    }
#pragma endregion

#pragma region start cipher
    if ((err = cbc_start(cipher_idx, IV, key, ks, 0, &cbc)) != CRYPT_OK)
    {
      printf("cbc_start error: %s\n", error_to_string(err));
      exit(-1);
    }
#pragma endregion

#pragma region figuer size of cipher text in octets
    /* we need to pad white space and amount of bytes we pad
        so that the plaintext is multiple of block length
      */
    octets /= ivsize;
    octets += 1;
#pragma endregion

#pragma endregion

#pragma region encryption loop
    for (long long i = 0; i < octets; i++)
    {
      y = fread(inbuf, 1, ivsize, fdin);

#pragma region padding
      if (i == octets - 1)
      {
        for (int cursor = y; cursor < sizeof(inbuf); cursor++)
        {
          inbuf[cursor] = ' ';
        }
        inbuf[ivsize - 1] = (unsigned char)y;
        y = ivsize;
        // printf("encrypt%s\n", inbuf);
      }
#pragma endregion

      if ((err = cbc_encrypt(inbuf, ciphertext, y, &cbc)) != CRYPT_OK)
      {
        printf("cbc_encrypt error: %s\n", error_to_string(err));
        exit(-1);
      }

      if (fwrite(ciphertext, 1, y, fdout) != y)
      {
        printf("Error writing to output.\n");
        exit(-1);
      }
    }
#pragma endregion

#pragma region release file handles
    fclose(fdout);
    fclose(fdin);
#pragma endregion
  }
#pragma endregion
}

void cfbmode()
{
#pragma region decryption
  if (decrypt)
  {

#pragma region initialization of algorithm

#pragma region read iv
    /* we need to read iv */
    if (fread(IV, 1, ivsize, fdin) != ivsize)
    {
      printf("Error reading IV from input.\n");
      exit(-1);
    }
#pragma endregion

#pragma region start cipher
    if ((err = cfb_start(cipher_idx, IV, key, ks, 0, &cfb)) != CRYPT_OK)
    {
      printf("cfb_start error: %s\n", error_to_string(err));
      exit(-1);
    }
#pragma endregion

#pragma region figuer size of plain text in octets
    /* we need to remove padding so that the plaintext is the same as origin */
    octets /= ivsize;
    octets++;
#pragma endregion

#pragma endregion

#pragma region decryption loop
    for (long long i = 1; i < octets; i++)
    {
      y = fread(inbuf, 1, ivsize, fdin);

      if ((err = cfb_decrypt(inbuf, plaintext, y, &cfb)) != CRYPT_OK)
      {
        printf("cfb_decrypt error: %s\n", error_to_string(err));
        exit(-1);
      }

      if (fwrite(plaintext, 1, y, fdout) != y)
      {
        printf("Error writing to file.\n");
        exit(-1);
      }
    }
#pragma endregion

#pragma region release file handles
    fclose(fdin);
    fclose(fdout);
#pragma endregion
  }
#pragma endregion

  else

#pragma region encryption
  { /* encrypt */

#pragma region initialization of algorithm

#pragma region write iv
    if ((err = rng_make_prng(128, find_prng("yarrow"), &prng, NULL)) != CRYPT_OK)
    {
      printf("Error setting up PRNG, %s\n", error_to_string(err));
    }

    /* You can use rng_get_bytes on platforms that support it */
    /* x = rng_get_bytes(IV,ivsize,NULL);*/
    x = yarrow_read(IV, ivsize, &prng);
    if (x != ivsize)
    {
      printf("Error reading PRNG for IV required.\n");
      exit(-1);
    }

    if (fwrite(IV, 1, ivsize, fdout) != ivsize)
    {
      printf("Error writing IV to output.\n");
      exit(-1);
    }
#pragma endregion

#pragma region start cipher
    if ((err = cfb_start(cipher_idx, IV, key, ks, 0, &cfb)) != CRYPT_OK)
    {
      printf("cfb_start error: %s\n", error_to_string(err));
      exit(-1);
    }
#pragma endregion

#pragma region figuer size of cipher text in octets
    /* we need to pad white space and amount of bytes we pad
        so that the plaintext is multiple of block length
      */
    octets /= ivsize;
    octets += 1;
#pragma endregion

#pragma endregion

#pragma region encryption loop
    for (long long i = 0; i < octets; i++)
    {
      y = fread(inbuf, 1, ivsize, fdin);

      if ((err = cfb_encrypt(inbuf, ciphertext, y, &cfb)) != CRYPT_OK)
      {
        printf("cfb_encrypt error: %s\n", error_to_string(err));
        exit(-1);
      }

      if (fwrite(ciphertext, 1, y, fdout) != y)
      {
        printf("Error writing to output.\n");
        exit(-1);
      }
    }
#pragma endregion

#pragma region release file handles
    fclose(fdout);
    fclose(fdin);
#pragma endregion
  }
#pragma endregion
}

void ofbmode()
{
#pragma region decryption
  if (decrypt)
  {

#pragma region initialization of algorithm

#pragma region read iv
    /* we need to read iv */
    if (fread(IV, 1, ivsize, fdin) != ivsize)
    {
      printf("Error reading IV from input.\n");
      exit(-1);
    }
#pragma endregion

#pragma region start cipher
    if ((err = ofb_start(cipher_idx, IV, key, ks, 0, &ofb)) != CRYPT_OK)
    {
      printf("ofb_start error: %s\n", error_to_string(err));
      exit(-1);
    }
#pragma endregion

#pragma region figuer size of plain text in octets
    /* we need to remove padding so that the plaintext is the same as origin */
    octets /= ivsize;
    octets++;
#pragma endregion

#pragma endregion

#pragma region decryption loop
    for (long long i = 1; i < octets; i++)
    {
      y = fread(inbuf, 1, ivsize, fdin);

      if ((err = ofb_decrypt(inbuf, plaintext, y, &ofb)) != CRYPT_OK)
      {
        printf("ofb_decrypt error: %s\n", error_to_string(err));
        exit(-1);
      }

      if (fwrite(plaintext, 1, y, fdout) != y)
      {
        printf("Error writing to file.\n");
        exit(-1);
      }
    }
#pragma endregion

#pragma region release file handles
    fclose(fdin);
    fclose(fdout);
#pragma endregion
  }
#pragma endregion

  else

#pragma region encryption
  { /* encrypt */

#pragma region initialization of algorithm

#pragma region write iv
    if ((err = rng_make_prng(128, find_prng("yarrow"), &prng, NULL)) != CRYPT_OK)
    {
      printf("Error setting up PRNG, %s\n", error_to_string(err));
    }

    /* You can use rng_get_bytes on platforms that support it */
    /* x = rng_get_bytes(IV,ivsize,NULL);*/
    x = yarrow_read(IV, ivsize, &prng);
    if (x != ivsize)
    {
      printf("Error reading PRNG for IV required.\n");
      exit(-1);
    }

    if (fwrite(IV, 1, ivsize, fdout) != ivsize)
    {
      printf("Error writing IV to output.\n");
      exit(-1);
    }
#pragma endregion

#pragma region start cipher
    if ((err = ofb_start(cipher_idx, IV, key, ks, 0, &ofb)) != CRYPT_OK)
    {
      printf("ofb_start error: %s\n", error_to_string(err));
      exit(-1);
    }
#pragma endregion

#pragma region figuer size of cipher text in octets
    /* we need to pad white space and amount of bytes we pad
        so that the plaintext is multiple of block length
      */
    octets /= ivsize;
    octets += 1;
#pragma endregion

#pragma endregion

#pragma region encryption loop
    for (long long i = 0; i < octets; i++)
    {
      y = fread(inbuf, 1, ivsize, fdin);

      if ((err = ofb_encrypt(inbuf, ciphertext, y, &ofb)) != CRYPT_OK)
      {
        printf("ofb_encrypt error: %s\n", error_to_string(err));
        exit(-1);
      }

      if (fwrite(ciphertext, 1, y, fdout) != y)
      {
        printf("Error writing to output.\n");
        exit(-1);
      }
    }
#pragma endregion

#pragma region release file handles
    fclose(fdout);
    fclose(fdin);
#pragma endregion
  }
#pragma endregion
}

void rc4cipher()
{
  rc4_state st;
  if ((err = rc4_stream_setup(&st, key, outlen)) != CRYPT_OK)
  {
    printf("rc4_stream_setup error: %s\n", error_to_string(err));
    exit(-1);
  }
  octets /= 128;
  octets++;
  for (long long i = 0; i < octets; i++)
  {
    y = fread(inbuf, 1, sizeof(inbuf), fdin);
    if ((err = rc4_stream_crypt(&st, inbuf, y, outbuf)) != CRYPT_OK)
    {
      printf("rc4_stream_crypt error: %s\n", error_to_string(err));
      exit(-1);
    }
    if (fwrite(outbuf, 1, y, fdout) != y)
    {
      printf("Error writing to output.\n");
      exit(-1);
    }
  }

#pragma region release file handles
  fclose(fdout);
  fclose(fdin);
#pragma endregion
}
#pragma endregion

int main(int argc, char *argv[])
{
#pragma region handle arguments
  if (argc < 4)
  {
    return usage(argv[0]);
  }

  if (strcmp(argv[1], "-d") == 0)
  {
    decrypt = 1;
    cipher = argv[2];
    mode = argv[3];
    infile = argv[4];
    outfile = argv[5];
  }
  else
  {
    decrypt = 0;
    cipher = argv[1];
    mode = argv[2];
    infile = argv[3];
    outfile = argv[4];
  }
#pragma endregion

#pragma region initialization

#pragma region register algorithms
  register_all_ciphers();
  register_all_hashes();
  register_all_prngs();
#pragma endregion

#pragma region setup cipher index

  cipher_idx = find_cipher(cipher);
  if (cipher_idx == -1 && strcmp(cipher, "rc4") != 0)
  {
    printf("Invalid cipher entered on command line.\n");
    exit(-1);
  }

#pragma endregion

#pragma region get the size of infile
  stat(infile, &file_info);
  octets = file_info.st_size;
#pragma endregion

#pragma region setup file handles
  fdin = fopen(infile, "rb");
  if (fdin == NULL)
  {
    perror("Can't open input for reading");
    exit(-1);
  }

  fdout = fopen(outfile, "wb");
  if (fdout == NULL)
  {
    printf("out err");
    perror("Can't open output for writing");
    exit(-1);
  }
#pragma endregion

#pragma region setup hash
  hash_idx = find_hash("sha256");
  if (hash_idx == -1)
  {
    printf("LTC_SHA256 not found...?\n");
    exit(-1);
  }
#pragma endregion

#pragma region setup iv
  ks = hash_descriptor[hash_idx].hashsize;
  if (strcmp(cipher, "rc4") != 0)
  {
    ivsize = cipher_descriptor[cipher_idx].block_length;

    if (cipher_descriptor[cipher_idx].keysize(&ks) != CRYPT_OK)
    {
      printf("Invalid keysize???\n");
      exit(-1);
    }
  }

#pragma endregion

#pragma region key generation
  printf("Enter key: ");
  if (fgets((char *)tmpkey, sizeof(tmpkey), stdin) == NULL)
  {
    exit(-1);
  }
  printf("\n");
  outlen = sizeof(key);
  if ((err = hash_memory(hash_idx, tmpkey, strlen((char *)tmpkey), key, &outlen)) != CRYPT_OK)
  {
    printf("Error hashing key: %s\n", error_to_string(err));
    exit(-1);
  }
#pragma endregion

#pragma endregion

#pragma region work area
  if (strcmp(cipher, "rc4") == 0)
  {
    rc4cipher();
  }
  else if (strcmp(mode, "ecb") == 0)
  {
    ecbmode();
  }
  else if (strcmp(mode, "cbc") == 0)
  {
    cbcmode();
  }
  else if (strcmp(mode, "cfb") == 0)
  {
    cfbmode();
  }
  else if (strcmp(mode, "ofb") == 0)
  {
    ofbmode();
  }
#pragma endregion

  return 0;
}